from hep.aida import *

from java.awt import Color
from java.awt import Font
from java.lang import *
from java.util import *
from java.io import *

from jhplot.jadraw import *
from jhplot  import *

from fr.in2p3.ipnl.aidautilities import *
from fr.in2p3.ipnl.oflap.scint import *
from jscinttools.utilities.jhplotutilities import *

Locale.setDefault(Locale.ENGLISH)
System.getProperties().setProperty("hep.aida.BatchAnalysisFactory","hep.aida.ref.BatchAnalysisFactory")
af = IAnalysisFactory.create();

showFigures = True
writeFigures = not showFigures

""" Define colors and fonts """
figType = [".png", ".eps"]
colors = [Color.blue, Color.red, Color(0, 170, 0), Color.magenta, Color.black, Color.yellow]
axisLabelFont = Font(Font.SERIF, Font.PLAIN, 25)
titleFont = Font(Font.SERIF, Font.PLAIN, 25)
legendFont = Font(Font.SERIF, Font.PLAIN, 20)
ticFont = Font(Font.SERIF, Font.PLAIN, 20)
statBoxFont = Font(Font.SERIF, Font.PLAIN, 15)

""" Fit Parameters """
k1_us = 123.9
k2_us = 2.5
d_K = 69.3
capK_us = 0.0067
deltaE_K = 1229.0
#n1 = 1.0
#kag1_us = 354.6
#kag2_us = 1.21
#dag_K = 103.3
#kag1_us = 355.6
#kag2_us = 5.3
#dag_K = 76.7

""" 3 lvl AG Fit
kag1_us = 356.0
kag2_us = 9.58
dag_K = 66.4
capKag_us = 0.000019
deltaEag_K = 3272.3
"""

""" 2 lvl AG Fit """
kag1_us = 350.0
kag2_us = 0.1
dag_K = 150



#kag1_us = 356.0
#kag2_us = 0.01
##dag_K = 66.4
#dag_K = 165

#capKag_us = 0.000019
#deltaEag_K = 3272.3
#n3 = 0.00001

#norm = 3.87
#n3Rel = 0.51

#norm = 50
#n3Rel = 0.51

#norm = 13.0
##n3Rel = 0.167
#n3Rel = 2.5

#norm = 4.0
#n3Rel = 0.167
#n3Rel = -0.2

norm = 10.0
n3Rel = 0.1

#norm = 3.691
#n3Rel = 0.0987

#lyPar = [k1_us, k2_us, d_K, capK_us, deltaE_K, kag1_us, norm, n3Rel]

lyPar = [k1_us, k2_us, d_K, capK_us, deltaE_K, kag1_us, kag2_us, dag_K, norm, n3Rel]

#lyPar = [k1_us, k2_us, d_K, capK_us, deltaE_K, kag1_us, kag2_us, dag_K, capKag_us, deltaEag_K, norm, n3Rel]
#lyPar = [k1_us, k2_us, d_K, capK_us, deltaE_K, norm]

""" Define the Path and files names """
""" Lyon """
aidaPath = "/home/marc-antoine/Documents/BGO/QFin"
aidaFile = "ScintillationAnalysis_BGO"
lyName = "/Light_Yield/LY Corr"

""" Queens """
#aidaPath = "/home/marc-antoine/Documents/BGO/QFin/QueensData/queenstools"
#aidaFile = "ScintillationAnalysis_BGO_QFile"
#lyName = "/Light_Yield/LY Corr Q"

outPath = "/home/marc-antoine/Dropbox/BGO"
tree = AidaFileReader(aidaPath, aidaFile).tree

""" Create the DPS with the 2 initials DPS and do a fit """
lyDPS = GrabIObject(tree, lyName).dps
aUtil = AidaUtil()
aUtil.sortUp(lyDPS)

fitter = af.createFitFactory().createFitter("Chi2", "jminuit");

"""
fitFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("LY Fit 3 Lvl AG", 1, "A*((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0]))*((1.0/((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0])+(1+exp(-1.0*D_K/x[0]))*(1.0/K_us)*exp(-1.0*DeltaE_K/x[0])))-(B/((1.0/kag1_us))))", "k1_us,k2_us,D_K,K_us,DeltaE_K,kag1_us,A,B", "LY Fit 3 Lvl AG")
fitFunc.setParameter("k1_us", lyPar[0])
fitFunc.setParameter("k2_us", lyPar[1])
fitFunc.setParameter("D_K", lyPar[2])
fitFunc.setParameter("K_us", lyPar[3])
fitFunc.setParameter("DeltaE_K", lyPar[4])
fitFunc.setParameter("kag1_us", lyPar[5])
fitFunc.setParameter("A", lyPar[6])
fitFunc.setParameter("B", lyPar[7])
"""


fitFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("LY Fit 3 Lvl AG", 1, "A*((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0]))*((1.0/((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0])+(1+exp(-1.0*D_K/x[0]))*(1.0/K_us)*exp(-1.0*DeltaE_K/x[0])))-(B/A/((1.0/kag1_us)+(1.0/kag2_us)*exp(-1.0*dEag_K/x[0]))))", "k1_us,k2_us,D_K,K_us,DeltaE_K,kag1_us,kag2_us,dEag_K,A,B", "LY Fit 3 Lvl AG")
fitFunc.setParameter("k1_us", lyPar[0])
fitFunc.setParameter("k2_us", lyPar[1])
fitFunc.setParameter("D_K", lyPar[2])
fitFunc.setParameter("K_us", lyPar[3])
fitFunc.setParameter("DeltaE_K", lyPar[4])
fitFunc.setParameter("kag1_us", lyPar[5])
fitFunc.setParameter("kag2_us", lyPar[6])
fitFunc.setParameter("dEag_K", lyPar[7])
fitFunc.setParameter("A", lyPar[8])
fitFunc.setParameter("B", lyPar[9])


"""
fitFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("LY Fit 3 Lvl AG", 1, "A*((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0]))*((1.0/((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0])+(1.0+exp(-1.0*D_K/x[0]))*(1.0/K_us)*exp(-1.0*DeltaE_K/x[0])))-(B/A/((1.0/kag1_us)+(1.0/kag2_us)*exp(-1.0*dEag_K/x[0])+(1.0+exp(-1.0*dEag_K/x[0]))*(1.0/Kag_us)*exp(-1.0*DeltaEag_K/x[0]))))", "k1_us,k2_us,D_K,K_us,DeltaE_K,kag1_us,kag2_us,dEag_K,Kag_us,DeltaEag_K,A,B", "LY Fit 3 Lvl AG")
fitFunc.setParameter("k1_us", lyPar[0])
fitFunc.setParameter("k2_us", lyPar[1])
fitFunc.setParameter("D_K", lyPar[2])
fitFunc.setParameter("K_us", lyPar[3])
fitFunc.setParameter("DeltaE_K", lyPar[4])
fitFunc.setParameter("kag1_us", lyPar[5])
fitFunc.setParameter("kag2_us", lyPar[6])
fitFunc.setParameter("dEag_K", lyPar[7])
fitFunc.setParameter("Kag_us", lyPar[8])
fitFunc.setParameter("DeltaEag_K", lyPar[9])
fitFunc.setParameter("A", lyPar[10])
fitFunc.setParameter("B", lyPar[11])
"""


doFit = True
fixPar = True

if fixPar:
	fitter.fitParameterSettings("k1_us").setFixed(True)
	fitter.fitParameterSettings("k2_us").setFixed(True)
	fitter.fitParameterSettings("D_K").setFixed(True)
	fitter.fitParameterSettings("K_us").setFixed(True)
	fitter.fitParameterSettings("DeltaE_K").setFixed(True)
	fitter.fitParameterSettings("kag1_us").setFixed(True)
	fitter.fitParameterSettings("kag2_us").setFixed(True)
	fitter.fitParameterSettings("dEag_K").setFixed(True)
	fitter.fitParameterSettings("Kag_us").setFixed(True)
	fitter.fitParameterSettings("DeltaEag_K").setFixed(True)
	#fitter.fitParameterSettings("A").setFixed(True)

#fitter.fitParameterSettings("B").setBounds(0.0,1.0) 

#fitter.fitParameterSettings("k1_us").setLowerBound(0.0) 
#fitter.fitParameterSettings("k2_us").setLowerBound(0.0) 
#fitter.fitParameterSettings("D_K").setLowerBound(0.0) 
#fitter.fitParameterSettings("K_us").setLowerBound(0.0) 
#fitter.fitParameterSettings("DeltaE_K").setLowerBound(0.0) 
#fitter.fitParameterSettings("kag1_us").setLowerBound(0.0) 
#fitter.fitParameterSettings("kag2_us").setLowerBound(0.0) 
#fitter.fitParameterSettings("dEag_K").setLowerBound(0.0) 
#fitter.fitParameterSettings("Kag_us").setLowerBound(0.0) 
#fitter.fitParameterSettings("DeltaEag_K").setLowerBound(0.0) 
#fitter.fitParameterSettings("A").setLowerBound(0.0)
fitter.fitParameterSettings("B").setLowerBound(0.0) 
#fitter.fitParameterSettings("A").setUpperBound(0)
#fitter.fitParameterSettings("B").setUpperBound(0)
func = fitFunc
if doFit:
	theFitResult = fitter.fit(lyDPS, fitFunc)
			
	print("\n================ Light Yield Fit with Three Level Model ================")
	print(AidaUtil.printResultsOfFit(theFitResult))
	func = theFitResult.fittedFunction()


#theFitResult = ScintDataFitter.doLYFit(lyDPS, lyPar)


"""
###############################################################
# Plotting Part
###############################################################
"""
""" Create the plottable JHep objects from AIDA IObjects """
""" The DataPointSet """
mainP1D = P1D(lyDPS)
mainP1D.setTitle("Main Component")
mainP1D.setStyle("p")
mainP1D.setSymbol(2)
mainP1D.setSymbolSize(12)
mainP1D.setColor(colors[3])
mainP1D.setPenWidth(2)
mainP1D.setErrY(True)
mainP1D.setErrX(False)
mainP1D.setErrColor(mainP1D.getColor())


""" The gamma fit fuction """
#mainF1D = F1D(theFitResult.fittedFunction(), 0.1, 1000)
mainF1D = F1D(func, 0.1, 1000)
mainF1D.setTitle("Fit &gamma;")
mainF1D.setPoints(10000)
mainF1D.setColor(colors[2])
mainF1D.setPenWidth(2)


mainP1DKey = JaKey(mainP1D.getSymbol(), mainP1D.getTitle(), 0.20, 0.65, "NDC")
mainP1DKey.setColor(mainP1D.getColor())
mainP1DKey.setKeyColor(mainP1D.getColor())
mainP1DKey.setPenWidth(2)
mainP1DKey.setFont(legendFont)


mainF1DKey = JaKey(20, mainF1D.getTitle(), 0.20, 0.75, "NDC")
mainF1DKey.setColor(mainF1D.getColor())
mainF1DKey.setKeyColor(mainF1D.getColor())
mainF1DKey.setPenWidth(2)
mainF1DKey.setFont(legendFont)


""" Create the Plotter """
mainPlotter = HPlotJa("LY", 600, 600, 1, 1, True)
mainPlotter.setGTitle("LY in BGO", titleFont)
mainPlotter.setAntiAlias(False)
if showFigures:
	mainPlotter.visible(True)
else:
	mainPlotter.visible(False)
mainPlotter.setPad(1, 1, 0.125, 0.125, 0.75, 0.75)
mainPlotter.setRange(1, 1000, 0.0, 7)
mainPlotter.setLogScale(0, True)
mainPlotter.setLogScale(1, False)
mainPlotter.setTicFont(ticFont)
mainPlotter.setShowStatBox(False)
mainPlotter.showKey(False)
mainPlotter.setNameX("Temperature (K)")
mainPlotter.setNameY("Light Yield (a.u.)")

""" Plots all the JHep objects """
mainPlotter.draw(mainP1D)
mainPlotter.draw(mainF1D)
#mainPlotter.draw(mainP1DL)
#mainPlotter.draw(mainP1DQ)
#mainPlotter.draw(mainP1DKeyL)
#mainPlotter.draw(mainP1DKeyQ)
#mainPlotter.draw(mainF1DKey)

""" Modifies the JaAxes and the axes labels """
labX = mainPlotter.getLabelX()[0][0]
labX.setFont(axisLabelFont)
HPlotJaUtil.setLabelLocationAtCenter(0, labX, mainPlotter)
labY = mainPlotter.getLabelY()[0][0]
labY.setFont(axisLabelFont)
HPlotJaUtil.setLabelLocationAtCenter(1, labY, mainPlotter)

theAxes = mainPlotter.getPad()
theAxes.setLabelSpace(0, 0.0)
theAxes.setLabelSpace(1, 0.0)
#theAxes.setLabelSpace(1, -0.025)
HPlotJaUtil.setTickLabelOneToPower(0, theAxes);

mainPlotter.update()

""" Saves or not the plot """
if writeFigures:
	for i in range(len(figType)):
		outFileName = "LYFit_3lvlAg" + figType[i]
		mainPlotter.export(outPath + File.separator + outFileName)
		print("Writting File : " + outPath + File.separator + outFileName)
		
mainPlotter.update()
print("That's all folks !")

